#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vector>
#include <map>
#include "pub.h"
#include "listen.h"
#include "nqlog.h"
#include "sysplat.h"

using namespace std;

IPMAP local_ip;
IPMAP remote_ip;
SVRMAP local_svr;
SVRMAP remote_svr;
SVRMAP dns_svr;
DOMAINMAP domain_list;

void NQ_write_localip(IPMAP& tmp)
{
    IPMAPITR it;
    local_ip.clear();
    local_ip = tmp;
    for (it = local_ip.begin(); it != local_ip.end(); it++)
    {
        NQResult rs;
        rs.nq_type = NQTYPE_LOCAL_IP;
        rs.nq_sip = it->first;        
        NQ_ret_write(rs);
    }
}

void NQ_write_localsvr(SVRMAP& tmp)
{
    SVRMAPITR it;
    local_svr.clear();
    local_svr = tmp;
    for (it = local_svr.begin(); it != local_svr.end(); it++)
    {
        NQResult rs;
        rs.nq_type = NQTYPE_LOCAL_SVR;
        rs.nq_sip = it->first.ip;
        rs.nq_sport = it->first.port;   
        DEBUG_OUT("%x:%u\n", it->first.ip, it->first.port);
        NQ_ret_write(rs);
    }
}

void NQ_remote_ip(UINT ip)
{
    IPMAPITR it_r;

    it_r = remote_ip.find(ip);
    if (remote_ip.end() == it_r)
    {
        NQResult rs;
        rs.nq_type = NQTYPE_REMOTE_IP;
        rs.nq_dip = ip;
        NQ_ret_write(rs);
        remote_ip.insert(make_pair(ip, ip));
    }
}

void NQ_update_ip(UINT sip, UINT dip)
{
    IPMAPITR it_l;
    IPMAPITR it_r;

    it_l = local_ip.find(sip);
    if (local_ip.end() == it_l)
    {
        NQ_remote_ip(sip);
    }
    else
    {
        NQ_remote_ip(dip);
    }
}

void NQ_remote_svr(UINT r_ip, UINT r_port)
{
    SVRMAPITR it_r;
    NQSVR_S nq_svr;

    nq_svr.ip = r_ip;
    nq_svr.port = ntohs(r_port);
    it_r = remote_svr.find(nq_svr);
    if (remote_svr.end() == it_r)
    {
        NQResult rs;
        rs.nq_type = NQTYPE_REMOTE_SVR;
        rs.nq_dip = nq_svr.ip;
        rs.nq_dport = nq_svr.port;
        NQ_ret_write(rs);
        remote_svr.insert(make_pair(nq_svr, nq_svr));
    }
}

void NQ_update_svr(NQResult& rs)
{
    IPMAPITR it_l;

    it_l = local_ip.find(rs.nq_sip);
    if (local_ip.end() == it_l)
    {
        NQ_remote_svr(rs.nq_sip, rs.nq_sport);
    }
}

void NQ_update_dns_svr(UINT svr_ip, UINT svr_port)
{
    SVRMAPITR it_r;
    NQSVR_S nq_svr;

    nq_svr.ip = svr_ip;
    nq_svr.port = ntohs(svr_port);
    it_r = dns_svr.find(nq_svr);
    if (dns_svr.end() == it_r)
    {
        NQResult rs;
        rs.nq_type = NQTYPE_DNS_SVR;
        rs.nq_dip = nq_svr.ip;
        rs.nq_dport = nq_svr.port;
        NQ_ret_write(rs);
        dns_svr.insert(make_pair(nq_svr, nq_svr));
    }
}

void NQ_update_domain(char *name, UINT ip)
{
    NQDOMAIN domain;
    DOMAINMAPITR it;

    domain.ip = ip;
    strcpy(domain.name, name);
    it = domain_list.find(domain);
    if (domain_list.end() == it)
    {
        NQResult rs;
        rs.nq_type = NQTYPE_DNS_NAME;
        rs.nq_dnsip = ip;
        strcpy(rs.nq_name, name);
        NQ_ret_write(rs);
        domain_list.insert(make_pair(domain, domain));
    }
}

void NQ_send_dns(UINT sip, UINT dip, char *name)
{
    NQResult rs;
    rs.nq_type = NQTYPE_SEND_DNS;
    //rs.nq_sip = sip;
    //rs.nq_dip = dip;
    strcpy(rs.nq_name, name);
    NQ_ret_write(rs);
}

void NQ_send_syn(NQResult& rs)
{
    IPMAPITR it_l;
    
    it_l = local_ip.find(rs.nq_sip);
    if (local_ip.end() != it_l)
    {
        rs.nq_type = NQTYPE_SEND_SYN;
        rs.nq_sport = ntohs(rs.nq_sport);
        rs.nq_dport = ntohs(rs.nq_dport);
        NQ_ret_write(rs);
    }
}

inline void NQ_ntoh_dnshdr(DNSHDR_S *dnshdr, DNSHDR_S *dnshdr_host)
{
    dnshdr_host->tid = ntohs(dnshdr->tid);
    dnshdr_host->flag = ntohs(dnshdr->flag);
    dnshdr_host->question_cnt = ntohs(dnshdr->question_cnt);
    dnshdr_host->answer_cnt = ntohs(dnshdr->answer_cnt);
    dnshdr_host->author_cnt = ntohs(dnshdr->author_cnt);
    dnshdr_host->addit_cnt = ntohs(dnshdr->addit_cnt);
}

inline void NQ_ntoh_dnsdata(DNSDATA_S *resp_data, DNSDATA_S *resp_data_host)
{
    resp_data_host->flag = resp_data->flag;
    resp_data_host->offset = resp_data->offset;
    resp_data_host->type = ntohs(resp_data->type);
    resp_data_host->qclass = ntohs(resp_data->qclass);
    resp_data_host->ttl = ntohl(resp_data->ttl);
    resp_data_host->data_len = ntohs(resp_data->data_len);
}

void NQ_listen_dns_req(IPHDR_S *iphdr, UCHAR *data, UINT len)
{
    DNSHDR_S *dnshdr;
    DNSHDR_S dnshdr_host;
    UCHAR *req_data;
    char name[256] = { 0 };
    UINT data_off = 0;
    UINT name_off = 0;
    UINT label_len = 0;
    UINT temp_off = 0;
    UINT req_cnt;
    UINT ans_cnt;
    char first_dns[256] = { 0 };
    UCHAR *dns_answer_start;
    UINT answer_off = 0;

    if (len < sizeof(DNSHDR_S))
    {
        return ;
    }

    dnshdr = (DNSHDR_S *)data;
    NQ_ntoh_dnshdr(dnshdr, &dnshdr_host);

    /* req-cnt */
    if (1 != dnshdr_host.question_cnt)
    {
        return ;
    }

    /* req */
    req_data = data + sizeof(DNSHDR_S);
    for (req_cnt = 0; req_cnt < dnshdr_host.question_cnt; req_cnt++)
    {
        memset(name, 0, sizeof(name));
        while (0 != req_data[data_off])
        {
            if (!NQ_DNSCHAR(req_data[data_off]))
            {
                label_len = req_data[data_off];
                temp_off = 0;
            }
            else
            {
                name[name_off] = req_data[data_off];
                name_off++;
                temp_off++;
                if ((temp_off == label_len) && 
                    (0 != req_data[data_off+1]))
                {
                    name[name_off] = '.';
                    name_off++;
                }
            }
            data_off++;
        }

        if ((0 == first_dns[0]) && (0 != name[0]))
        {
            strcpy(first_dns, name);
            DEBUG_OUT("___DNS_REQ___: first name = %s\n", first_dns);
            NQ_send_dns(iphdr->sip, iphdr->dip, first_dns);
        }
        else
        {
            DEBUG_OUT("___DNS_REQ___: other name = %s\n", name);
        }
    }

    /* answer */
    dns_answer_start = req_data + data_off + 5;
    for (ans_cnt = 0; ans_cnt < dnshdr_host.answer_cnt; ans_cnt++)
    {
        UINT ip;
        DNSDATA_S *resp_data;
        DNSDATA_S resp_data_host;

        resp_data = (DNSDATA_S *)(dns_answer_start + answer_off);
        NQ_ntoh_dnsdata(resp_data, &resp_data_host);

        if ((0x01 != resp_data_host.type) || 
            (0x01 != resp_data_host.qclass) ||
            (4 != resp_data_host.data_len))
        {
            answer_off = answer_off + (sizeof(DNSDATA_S) + resp_data_host.data_len);
            continue;
        }

        /* ip */
        ip = *((UINT *)(resp_data + sizeof(DNSDATA_S)));
        DEBUG_OUT("___DNS_REQ___: rest: first name = %s, ip = %s\n", first_dns, NQ_IPSTR(ip));
    }
}

void NQ_listen_dns_resp(IPHDR_S *iphdr, UCHAR *data, UINT len)
{
    DNSHDR_S *dnshdr;
    DNSHDR_S dnshdr_host;
    UCHAR *req_data;
    char name[256] = { 0 };
    UINT data_off = 0;
    UINT name_off = 0;
    UINT label_len = 0;
    UINT temp_off = 0;
    UINT req_cnt;
    UINT ans_cnt;
    char first_dns[256] = { 0 };
    UCHAR *dns_answer_start;
    UINT answer_off = 0;

    if (len < sizeof(DNSHDR_S))
    {
        return ;
    }

    dnshdr = (DNSHDR_S *)data;
    NQ_ntoh_dnshdr(dnshdr, &dnshdr_host);

    /* req-cnt */
    if (1 != dnshdr_host.question_cnt)
    {
        return ;
    }

    /* req */
    req_data = data + sizeof(DNSHDR_S);
    for (req_cnt = 0; req_cnt < dnshdr_host.question_cnt; req_cnt++)
    {
        memset(name, 0, sizeof(name));
        while (0 != req_data[data_off])
        {
            if (!NQ_DNSCHAR(req_data[data_off]))
            {
                label_len = req_data[data_off];
                temp_off = 0;
            }
            else
            {
                name[name_off] = req_data[data_off];
                name_off++;
                temp_off++;
                if ((temp_off == label_len) && 
                    (0 != req_data[data_off+1]))
                {
                    name[name_off] = '.';
                    name_off++;
                }
            }
            data_off++;
        }

        if ((0 == first_dns[0]) && (0 != name[0]))
        {
            strcpy(first_dns, name);
            DEBUG_OUT("___DNS_RESP___: first name = %s\n", first_dns);
        }
        else
        {
            DEBUG_OUT("___DNS_RESP___: other name = %s\n", name);
        }
    }

    /* answer */
    dns_answer_start = req_data + data_off + 5;
    for (ans_cnt = 0; ans_cnt < dnshdr_host.answer_cnt; ans_cnt++)
    {
        UINT ip;
        DNSDATA_S *resp_data;
        DNSDATA_S resp_data_host;

        resp_data = (DNSDATA_S *)(dns_answer_start + answer_off);
        NQ_ntoh_dnsdata(resp_data, &resp_data_host);
        
        if ((0x01 != resp_data_host.type) || 
            (0x01 != resp_data_host.qclass) ||
            (4 != resp_data_host.data_len))
        {
            answer_off = answer_off + (sizeof(DNSDATA_S) + resp_data_host.data_len);
            continue;
        }

        /* ip */
        ip = *((UINT *)(resp_data + 1));
        DEBUG_OUT("___DNS_RESP___: first name = %s, ip = %s\n", first_dns, NQ_IPSTR(ip));
        answer_off = answer_off + (sizeof(DNSDATA_S) + resp_data_host.data_len);

        /* update */
        NQ_update_domain(first_dns, ip);
    }
}

void NQ_listen_tcp(IPHDR_S *iphdr, UCHAR *data, UINT len)
{
    TCPHDR_S *tcphdr;
    UINT tcphdr_len;
    UINT payload_len;
    UCHAR *payload;

    if (len < sizeof(TCPHDR_S))
    {
        return ;
    }

    tcphdr = (TCPHDR_S *)data;
    if (FLAG_TCP_SYNACK == tcphdr->flag)
    {
        /* inet_ntoa bug */
        NQResult rs;
        rs.nq_sip = iphdr->sip;
        rs.nq_dip = iphdr->dip;
        rs.nq_sport = tcphdr->sport;
        rs.nq_dport = tcphdr->dport;
        NQ_update_svr(rs);
        DEBUG_OUT("[TCP-RESP] %s:%u ", NQ_IPSTR(iphdr->dip), ntohs(tcphdr->dport));     
        DEBUG_OUT("<- %s:%u\n", NQ_IPSTR(iphdr->sip), ntohs(tcphdr->sport));
    }
    else if (FLAG_TCP_SYN == tcphdr->flag)
    {
        NQResult rs;
        rs.nq_sip = iphdr->sip;
        rs.nq_dip = iphdr->dip;
        rs.nq_sport = tcphdr->sport;
        rs.nq_dport = tcphdr->dport;
        NQ_send_syn(rs);
        DEBUG_OUT("[TCP-REQ] %s:%u ", NQ_IPSTR(iphdr->sip), ntohs(tcphdr->sport));     
        DEBUG_OUT("-> %s:%u\n", NQ_IPSTR(iphdr->dip), ntohs(tcphdr->dport));
    }

    tcphdr_len = NQ_TCPHDR_LEN(tcphdr->len);
    if (len <= tcphdr_len)
    {
        return ;
    }
    
    payload = data + tcphdr_len;
    payload_len = len - tcphdr_len;
    
#if 0
    if (NQ_ISDNS(tcphdr->dport)) 
    {
        DEBUG_OUT("___TCP_DNS_REQ___\n");
        NQ_listen_dns_req(iphdr, payload, payload_len);
    }
#endif

    if (NQ_ISDNS(tcphdr->sport)) 
    {
        DEBUG_OUT("___TCP_DNS_RESP___\n");
        NQ_listen_dns_resp(iphdr, payload, payload_len);
    }

    return ;
}

void NQ_listen_udp(IPHDR_S *iphdr, UCHAR *data, UINT len)
{
    UDPHDR_S *udphdr;
    UINT payload_len;
    UCHAR *payload;

    if (len <= sizeof(UDPHDR_S))
    {
        return ;
    }

    udphdr = (UDPHDR_S *)data;
    if (len > ntohs(udphdr->len))
    {
        return ;
    }

    payload = data + sizeof(UDPHDR_S);
    payload_len = len - sizeof(UDPHDR_S);

    if (NQ_ISDNS(udphdr->dport)) 
    {
        NQ_listen_dns_req(iphdr, payload, payload_len);
    }
    
    if (NQ_ISDNS(udphdr->sport)) 
    {
        NQ_update_dns_svr(iphdr->sip, udphdr->sport);
        NQ_listen_dns_resp(iphdr, payload, payload_len);
    }
    
    return ;
}

void NQ_listen_icmp(IPHDR_S *iphdr, UCHAR *data, UINT len)
{
}

int NQ_listen_ip(UCHAR *data, UINT len)
{
    UINT iphdr_len;
    IPHDR_S *iphdr;
    UINT payload_len;
    UCHAR *payload;

    iphdr = (IPHDR_S *)data;

    if (!ISUNICAST(iphdr->sip) || !ISUNICAST(iphdr->dip))
    {
        return 0;
    }

    /* ignore loopback */
    if (iphdr->sip == iphdr->dip)
    {
        return 0;
    }

    /* update remote ip */
    NQ_update_ip(iphdr->sip, iphdr->dip);

    iphdr_len = 4 * (iphdr->ver & 0x0F);
    if (len <= iphdr_len)
    {
        return 0;
    }
    
    payload = data + iphdr_len;
    payload_len = len - iphdr_len;
    
    switch(iphdr->proto)
    {
        case IPPROTO_TCP:
            //DEBUG_OUT("___TCP___\n");
            NQ_listen_tcp(iphdr, payload, payload_len);
            break;
        case IPPROTO_UDP:
            //DEBUG_OUT("___UDP___\n");
            NQ_listen_udp(iphdr, payload, payload_len);
            break;
        case IPPROTO_ICMP:
            //DEBUG_OUT("___ICMP___\n");
            //NQ_listen_icmp(iphdr, payload, payload_len);
            break;
        default:
            break;
    }

    return 0;
}

inline void NQ_ntoh_arphdr(ARPHDR_S *arp)
{
    arp->h_type = ntohs(arp->h_type);
    arp->p_type = ntohs(arp->p_type);
    arp->opt    = ntohs(arp->opt);
}

int NQ_listen_arp(UCHAR *data, UINT len)
{
    ARPHDR_S *arp = (ARPHDR_S *)data;

    //DEBUG_OUT("___ARP___\n");  
    if (len < sizeof(ARPHDR_S))
    {
        return 0;
    }

    NQ_ntoh_arphdr(arp);

    /* Only listen ARP-response */
    if ((arp->h_type != 1) || (arp->p_type != 0x0800) ||
        (arp->h_len  != 6) || (arp->p_len  != 4) || (arp->opt != 2))
    {
        return 0;
    }

    NQ_update_ip(arp->sip, arp->dip);

    return 0;
}
